{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3622c575",
   "metadata": {
    "id": "3622c575"
   },
   "source": [
    "# TinyLlama Agent Demo\n",
    "\n",
    "This notebook demonstrates how to create a simple AI agent using the TinyLlama (1.1B) model.\n",
    "\n",
    "**Requirements:**\n",
    "- `transformers`\n",
    "- `accelerate`\n",
    "- `torch`\n",
    "\n",
    "We will load the model, prepare a basic conversational interface, and test it with a simple prompt."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "qS0M9pRhTGpf",
   "metadata": {
    "id": "qS0M9pRhTGpf"
   },
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/DhivyaBharathy-web/PraisonAI/blob/main/examples/cookbooks/TinyLlama_1_1B_model_SimpleAIAgent.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tuw0IDpBS0ei",
   "metadata": {
    "id": "tuw0IDpBS0ei"
   },
   "source": [
    "# Install necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0d06cae9",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0d06cae9",
    "outputId": "3e1d8fb7-19be-455b-dc35-d0ab19146669"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m363.4/363.4 MB\u001b[0m \u001b[31m3.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m13.8/13.8 MB\u001b[0m \u001b[31m29.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m24.6/24.6 MB\u001b[0m \u001b[31m29.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m883.7/883.7 kB\u001b[0m \u001b[31m20.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m664.8/664.8 MB\u001b[0m \u001b[31m2.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m211.5/211.5 MB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m12.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m127.9/127.9 MB\u001b[0m \u001b[31m7.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m207.5/207.5 MB\u001b[0m \u001b[31m6.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m45.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h"
     ]
    }
   ],
   "source": [
    "!pip install -q transformers accelerate torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "RFhYLkO1S5F9",
   "metadata": {
    "id": "RFhYLkO1S5F9"
   },
   "source": [
    "# Load TinyLlama model and tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "415af398",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 241,
     "referenced_widgets": [
      "dae7d635257948149c74e65956f614e3",
      "17ec6fcc9f624cdbb62fe2b8b55230d0",
      "58474504488044138a0252bb75258d86",
      "04fe4518fac04b9f8b29c62ed70f1dfd",
      "c6a9d3ea25594170977243d2048d3408",
      "730bebf00be142d3979c6d12b55d6757",
      "03ad1674c0884e7b9c4a21af0a1e66c3",
      "ebac47c887f04b79a60bd0918002bf0d",
      "ece9bce5d2ca4dd7b019040175513118",
      "458f11dbadb045f1aa8c72780e7aa2a0",
      "ec28a20d0a14445489723684d2d69c0a",
      "d347fa0c0fc849559b66c8f116c86317",
      "117f649fa62b4f5d9add0160e6cd80a4",
      "66f702418fa348eea27778282dc443aa",
      "44413339df5b4ae688f8101faf9ed3b5",
      "00cc576df95441c8916ce38a7c6decc2",
      "22e9ea6d191841bf9a5a5072326aa6d8",
      "bd94a291b3164113b5ea460b38b8c167",
      "9d132a297f924a3d867605df81e70d8e",
      "5345feddca204db784ea0452e739eda6",
      "ace33e6e88d842c89efd180b3a83f674",
      "61dd85a7edf544f196c60a259a91fa0d",
      "7451ec9bd091441c852d075b47bd1ca5",
      "05e3446e7c9243819a25d1ca8e8fa38e",
      "73141c4c9218462da97999d3afa459bc",
      "1cffc33252b94fbc861edb0953f13153",
      "de917c0fc49040c8a9254333c34c884a",
      "50962ececc0b42c28170eb0c23fd6171",
      "856cc834a33a4f9fb59c77f748006b36",
      "17fd1d476c1f4891b8cbeb8cacdeb866",
      "c97dc50f42cb44b8a71aeff5272d5337",
      "2e3a66f622d445c4be544a9306798cb6",
      "92920b1d428a4ed68a9e0329e2034cad",
      "7e282c1fcae842cc87c93ecfc1763654",
      "f9ed70f887204e2fb336e38a35096eec",
      "a4d99f9beca44e4a8fadff65b27aa78f",
      "b4f77b96d00445eea77ac78baa8c1e99",
      "695898e4a4c8476cb318446a6da592ba",
      "d70bcacbe77c41f685a4f8020f767889",
      "af49afc49ecf47fb9ebf2d299954eff2",
      "1bd2a253348f4b45928c875940f72481",
      "b6c5f55fc85b4331a2e749d27ce15f89",
      "81a87638c71046d7aeb1c642558602f7",
      "e750883103444391bd64872344c651d7",
      "e78dca0cf1314d01a5d5126e448e9d87",
      "d0c53dffc0c34d7090febf78231972b2",
      "152a118177914c6692e7ee4987b92b9b",
      "88fd71fd6dbf43ebb0a3b8471fd35f8d",
      "ddb91266f8834e4eaaf136cc15f4a831",
      "9f8179d135054eb292c9adc7fd66a91a",
      "297d5bc8b3a04f26b80e2cc5ee3f7bbd",
      "f49f72679cb944cd8fce5557bbd68b1b",
      "2ee28c8ef4f74690b80a89050da4ebb5",
      "c29924d6561a4feaa9238d1b34394941",
      "03636146b7c14f5db80b8f105be55303",
      "b391081de4444859ac3213aac99cc33a",
      "2babb83499fc46dc8422473f7b98403c",
      "a32f6cc2777941bca5b6250e36624ead",
      "90cd9996184a452ab651f9b8cc590f47",
      "dc04209c92764096ba9bd4aefa81c189",
      "16e5b93c3c964a8a8ca460d9077d0638",
      "794f22ded07d46c5a916325109618987",
      "dce367f1a5a74ebf94aea37209aa0a40",
      "5e355e3f340f4a7590e4a88d44b03383",
      "4e725e19c877462aa3c10726a335bbde",
      "d60ab6c255d244098a459459016a2089",
      "f1b07183320d42af88f985bd1abf054f",
      "d23103bb4b354e29a2c3e7dc93867b36",
      "ea3e0ae0aef54e5da113c9686ebbbab9",
      "4d3365c1bb7740dc8f36f4b419af135a",
      "c06ccbc051264027bd2c6521fec69abe",
      "597dbf5b28134cc0bf9d613443243455",
      "faa4f285eb224bbcbed649bd479e571f",
      "6e684d35885a4706b8add6834cd521ad",
      "f9c55ada3386460183e02559defec8a4",
      "9d64cbe1605f4f358923754c5f8b6424",
      "6008ed7eca09485195eb37c36acfbe79"
     ]
    },
    "id": "415af398",
    "outputId": "28de1886-e83f-49bf-b05f-cd19439605c5"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dae7d635257948149c74e65956f614e3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/1.29k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d347fa0c0fc849559b66c8f116c86317",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.model:   0%|          | 0.00/500k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7451ec9bd091441c852d075b47bd1ca5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.84M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e282c1fcae842cc87c93ecfc1763654",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/551 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e78dca0cf1314d01a5d5126e448e9d87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/608 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b391081de4444859ac3213aac99cc33a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/2.20G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f1b07183320d42af88f985bd1abf054f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "\n",
    "model_name = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "N0bRTd_DS8aL",
   "metadata": {
    "id": "N0bRTd_DS8aL"
   },
   "source": [
    "# Define a function to generate responses from the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1b0483e1",
   "metadata": {
    "id": "1b0483e1"
   },
   "outputs": [],
   "source": [
    "def generate_response(prompt, max_length=256):\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
    "    outputs = model.generate(**inputs, max_new_tokens=max_length)\n",
    "    return tokenizer.decode(outputs[0], skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vbHKNxc-S_Bk",
   "metadata": {
    "id": "vbHKNxc-S_Bk"
   },
   "source": [
    "# Example usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "298c9811",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "298c9811",
    "outputId": "6e412608-735c-4604-8a5d-a66b98dd9df0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a helpful assistant. What is the capital of France?\n"
     ]
    }
   ],
   "source": [
    "prompt = \"You are a helpful assistant. What is the capital of France?\"\n",
    "response = generate_response(prompt)\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}